from fastapi import FastAPI
import uvicorn
from pydantic import BaseModel
from server.inference import qwen2_5_14b_inference
from typing import List, Dict

app = FastAPI()

class Item(BaseModel):
    messages: List[Dict[str, str]]
    
# qwen_14b 推理路由
@app.on_event("startup")
@app.get("/qwen_14b_init_vllm")
def create_item():
    ret =  qwen2_5_14b_inference.qwen_14b_init_vllm()
    return {"message": f"{ret}"}


@app.post("/task/chat_qwen_14b")
def chat(item: Item):
    print('调用千问14b')
    ret = qwen2_5_14b_inference.qwen_14B_inference(item.messages)
    return ret

@app.post("/task/qwen_14b_judge")
def chat(item: Item):
    print('千问14b进行裁决')
    ret = qwen2_5_14b_inference.qwen_14B_judge(item.messages)
    return ret